package biz

import (
	"github.com/fullstorydev/grpcurl"
	"github.com/go-kratos/kratos/v2/registry"
	"google.golang.org/grpc"

	"gitee.com/shuokeyun/kratos/client"
	"gitee.com/shuokeyun/kratos/server"
)

type GrpcClientConnManage struct {
	registry registry.Discovery
	source   string
	conn     map[string]*grpc.ClientConn
}

func NewGrpcClientConnManage(app *server.App, source grpcurl.DescriptorSource) (*GrpcClientConnManage, func(), error) {
	m := &GrpcClientConnManage{
		registry: app.GetDiscovery(),
		source:   app.GetConfigBootstrap().GetServer().GetName(),
		conn:     make(map[string]*grpc.ClientConn),
	}
	serverList, err := source.ListServices()
	if err != nil {
		return nil, nil, err
	}
	err = m.initClientConn(serverList)
	if err != nil {
		return nil, nil, err
	}
	return m, func() {
		for _, v := range m.conn {
			v.Close()
		}
	}, nil
}

func (m *GrpcClientConnManage) GetClientConn(serviceName string) (*grpc.ClientConn, bool) {
	conn, ok := m.conn[serviceName]
	return conn, ok
}

func (m *GrpcClientConnManage) initClientConn(serverList []string) error {
	for _, serviceName := range serverList {
		conn, err := client.NewGrpcClientConn(m.registry, serviceName, m.source)
		if err != nil {
			return err
		}
		m.conn[serviceName] = conn
	}
	return nil
}
